package org.smartloli.kafka.eagle.infra.util;

import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.parser.CCJSqlParserManager;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.Limit;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectItem;
import net.sf.jsqlparser.util.TablesNamesFinder;
import org.apache.http.util.Asserts;
import org.springframework.util.CollectionUtils;

import java.io.StringReader;
import java.util.ArrayList;
import java.util.List;

/**
 * function:
 * author: zhiwei_yang
 * time: 2021/11/9-16:59
 */
@Slf4j
public final class JSqlParser {

    private JSqlParser() {
    }

    /**
     * 获取SQL查询字段
     *
     * @param sql
     * @return
     * @throws JSQLParserException
     */
    public static List<String> parseSelectItems(String sql) throws JSQLParserException {
        Asserts.notNull(sql, "Ksql is null");
        CCJSqlParserManager parserManager = new CCJSqlParserManager();
        Select select = (Select) parserManager.parse(new StringReader(sql));
        PlainSelect plain = (PlainSelect) select.getSelectBody();
        List<SelectItem> selectItems = plain.getSelectItems();
        List<String> selectItemNames = new ArrayList<>();
        if (!CollectionUtils.isEmpty(selectItems)) {
            for (SelectItem selectitem : selectItems) {
                selectItemNames.add(selectitem.toString());
            }
        }
        return selectItemNames;
    }

    /**
     * 判断是否是select语句
     *
     * @param sql
     * @return
     */
    public static boolean isSelect(String sql) {
        Asserts.notNull(sql, "Ksql is null");
        try {
            return CCJSqlParserUtil.parse(sql) instanceof Select;
        } catch (JSQLParserException e) {
            log.error("Ksql select method parse error ksql:{}", sql, e);
            return false;
        }
    }

    /**
     * 获取TopicName
     *
     * @param sql
     */
    public static String parseTopicName(String sql) {
        Asserts.notNull(sql, "Ksql is null");
        Statement statement = null;
        try {
            statement = CCJSqlParserUtil.parse(sql);
            Select selectStatement = (Select) statement;
            TablesNamesFinder tablesNamesFinder = new TablesNamesFinder();
            List<String> topicNames = tablesNamesFinder.getTableList(selectStatement);
            return CollectionUtils.isEmpty(topicNames) ? null : topicNames.get(0);
        } catch (JSQLParserException e) {
            log.error("Ksql topic name is empty, error ksql:{}", sql, e);
            return null;
        }
    }


    /**
     * 解析where条件: 只支持partition in(1，2，3) 和 partition = 1 场景
     */
    public static String parseWhereCondition(String sql) {
        Asserts.notNull(sql, "Ksql is null");
        try {
            CCJSqlParserManager parserManager = new CCJSqlParserManager();
            Select select = (Select) parserManager.parse(new StringReader(sql));
            PlainSelect plain = (PlainSelect) select.getSelectBody();
            Expression whereExpression = plain.getWhere();
            if (whereExpression instanceof InExpression || whereExpression instanceof EqualsTo) {
                return whereExpression.toString().toLowerCase();
            }
            return null;
        } catch (JSQLParserException e) {
            log.error("Ksql where condition is empty, error ksql:{}", sql, e);
            return null;
        }
    }

    /**
     * 解析limit条件
     */
    public static String parseLimitCondition(String sql) {
        Asserts.notNull(sql, "Ksql is null");
        try {
            CCJSqlParserManager parserManager = new CCJSqlParserManager();
            Select select = (Select) parserManager.parse(new StringReader(sql));
            PlainSelect plain = (PlainSelect) select.getSelectBody();
            Limit limitExpression = plain.getLimit();
            return null != limitExpression ? limitExpression.toString().toLowerCase() : null;
        } catch (JSQLParserException e) {
            log.error("Ksql limit condition is empty, error ksql:{}", sql, e);
            return null;
        }
    }

    public static void main(String[] args) throws JSQLParserException {
        String sql = "select name, passwd from user where partition in (1,2,3) limit 10";
        System.out.println(JSON.toJSONString(parseSelectItems(sql)));
        System.out.println(parseWhereCondition(sql));
        System.out.println(parseLimitCondition(sql));
    }
}